//////////////////////////////////////////////////////////////////////////////////
//	Copyright 2011 by Yosef Grabivker, ryuho@homenetworksoft.com
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//////////////////////////////////////////////////////////////////////////////////

#include "sshsocket.h"

#ifdef USE_SSH_SOCKET

namespace SSH2
{
#ifdef LINUX
	#include "../linux/libssh2_config.h"
#else
#include <libssh2_config.h>
#endif
#include <libssh2.h>
};

using namespace SSH2;

 

#ifdef WIN32
	#pragma warning (disable:4099)

	#ifdef _DEBUG
		#pragma comment(lib,"libssh2d.lib")
		#pragma comment(lib,"libeay32MDd.lib")
		#pragma comment(lib,"ssleay32MDd.lib")
	#else
		#pragma comment(lib,"libssh2.lib")
		#pragma comment(lib,"libeay32MD.lib")
		#pragma comment(lib,"ssleay32MD.lib")
	#endif
#endif

static std::string gstrPassword;

static void kbd_callback(const char *name, int name_len,
                         const char *instruction, int instruction_len,
                         int num_prompts,
                         const LIBSSH2_USERAUTH_KBDINT_PROMPT *prompts,
                         LIBSSH2_USERAUTH_KBDINT_RESPONSE *responses,
                         void **abstract)
{
    (void)name;
    (void)name_len;
    (void)instruction;
    (void)instruction_len;

    if (num_prompts == 1)
	{
        responses[0].text = strdup(gstrPassword.c_str());
        responses[0].length = gstrPassword.size();
    }

    (void)prompts;
    (void)abstract;
} /* kbd_callback */ 
 


CRSSHSocket::CRSSHSocket(int nTimeOut):CRTSocket(nTimeOut),
	m_pSession(NULL), m_pChannel(NULL)
{
	#if LIBSSH2_VERSION_NUM>=0x010206
	
	int rc;

	rc = libssh2_init (0);

    if (rc != 0)
	{
        fprintf (stderr, "SSH initialization failed (%d)\n", rc);
		THROW_EXC(cERFailed);       
    }
    
	#endif
}

CRSSHSocket::~CRSSHSocket(void)
{
}

void CRSSHSocket::Connect(const char* szAddress, unsigned short wPort)
{
	CRTSocket::Connect(szAddress, wPort);

	m_pSession = libssh2_session_init();

	if ( !m_pSession ) THROW_EXC(cERFailed);

    if ( 0!=libssh2_session_startup(m_pSession, GetSockHandle()) )
	{
        //fprintf(stderr, "Failure establishing SSH session\n");
        THROW_EXC(cERStartSSHConversation);
    }

	libssh2_hostkey_hash(m_pSession, LIBSSH2_HOSTKEY_HASH_SHA1);
}

void CRSSHSocket::Login( const CONPARAMS& con )
{
	int  auth_pw = 0;
	char *userauthlist;

	if ( !m_pSession ) THROW_EXC(cERFailed);

	RTSTRING strUser = (const char*)con.strUsr.ToAscii();
	RTSTRING strPasw = (const char*)con.strPassw.ToAscii();
	RTSTRING strPublickey = (const char*)con.strPublickey.ToAscii(); 
	RTSTRING strPrivkey = (const char*)con.strPrivkey.ToAscii();
	RTSTRING strPassphrase = (const char*)con.strPassphrase.ToAscii();

	userauthlist = libssh2_userauth_list(m_pSession, strUser.c_str(), strUser.size() );

	if ( !userauthlist ) THROW_EXC(cERFailed);

    //printf("Authentication methods: %s\n", userauthlist);

    if (strstr(userauthlist, "password") != NULL)
	{
        auth_pw |= 1;
    }
    if (strstr(userauthlist, "keyboard-interactive") != NULL)
	{
        auth_pw |= 2;
    }

    if (strstr(userauthlist, "publickey") != NULL)
	{
        auth_pw |= 4;
    }

	if ( con.bEnableSSHKey  )
	{
		if ( con.strPublickey.empty() ||
			 con.strPrivkey.empty() )
		{
			THROW_EXC(cERSSHKeysEmpty);
		}

		if ( auth_pw & 4 )
		{
			/* Or by public key */ 
			if ( libssh2_userauth_publickey_fromfile(m_pSession, strUser.c_str(),
				strPublickey.c_str(), strPrivkey.c_str(), strPassphrase.c_str()) )
			{
				//printf("\tAuthentication by public key failed!\n");
				THROW_EXC(cERSSHKeysAuthFailed);
			}
		}
		else
		{
			THROW_EXC(cERSSHKeysNotSupported);
		}
    }
	else
    if (auth_pw & 1)
	{
        /* We could authenticate via password */ 
		if ( libssh2_userauth_password(m_pSession, strUser.c_str(), strPasw.c_str()) )
		{
            //printf("\tAuthentication by password failed!\n");
			THROW_EXC(cERWrongUserOrPassword);
        }
		else
		{
            //printf("\tAuthentication by password succeeded.\n");
        }
    }
	else if (auth_pw & 2)
	{
		gstrPassword = strPasw;

        /* Or via keyboard-interactive */ 
        if (libssh2_userauth_keyboard_interactive(m_pSession, strUser.c_str(),
                                                  &kbd_callback) )
		{
            //printf("\tAuthentication by keyboard-interactive failed!\n");
            //goto shutdown;
		    THROW_EXC(cERWrongUserOrPassword);
        }
		else
		{
            //printf("\tAuthentication by keyboard-interactive succeeded.\n");
        }
    }
	else
	{
        //printf("No supported authentication methods found!\n");
        //goto shutdown;
		THROW_EXC(cERFailed);
    }

    /* Request a shell */ 
    if (!(m_pChannel = libssh2_channel_open_session(m_pSession)))
	{
        //fprintf(stderr, "Unable to open a session\n");
        THROW_EXC(cERFailed);
    }
 
    /* Some environment variables may be set,
     * It's up to the server which ones it'll allow though
     */ 
    //libssh2_channel_setenv(channel, "FOO", "bar");

 
    /* Request a terminal with 'vanilla' terminal emulation
     * See /etc/termcap for more options
     */ 
    if (libssh2_channel_request_pty(m_pChannel, "vanilla"))
	{
        //fprintf(stderr, "Failed requesting pty\n");
        THROW_EXC(cERFailed);
    }
 
    /* Open a SHELL on that pty */ 
    if (libssh2_channel_shell(m_pChannel))
	{
        //fprintf(stderr, "Unable to request shell on allocated pty\n");
        //goto shutdown;
		THROW_EXC(cERFailed);
    }
}

int CRSSHSocket::Write(const  char* pData, size_t stSize)
{
	int nRes(0);

	if (!m_pChannel) THROW_EXC(cERWrongObjectState);

	nRes = libssh2_channel_write(m_pChannel,pData, stSize);

	if ( nRes<0 )
	{
		CheckError();

		nRes = 0;
	}

	return nRes;
}

int CRSSHSocket::Read( char* pData, size_t stSize)
{
	int nRes(0);

	if (!m_pChannel) THROW_EXC(cERWrongObjectState);

	if ( WaitForRead(250) )
	{
		nRes = libssh2_channel_read(m_pChannel,pData, stSize);

		if ( nRes<0 )
		{
			CheckError();

			nRes = 0;
		}
	}

	return nRes;
}

void CRSSHSocket::Close()
{
	if (m_pChannel)
	{
		libssh2_channel_free(m_pChannel);
		m_pChannel = NULL;
	}

	if (m_pSession)
	{
		libssh2_session_disconnect(m_pSession, "Done");
		libssh2_session_free(m_pSession);

		m_pSession = NULL;
	}

	CRTSocket::Close();
}

void CRSSHSocket::Purge()
{
	char buf[1024];

	//must be while here
	while ( WaitForRead(25) )
	{
		if ( Read(buf, sizeof(buf))<=0 )
			break;
	}
}

#endif //USE_SSH_SOCKET
